# coding=UTF-8
import random
print("start to build train data pair!")
from const_ import *
#with open("data",encoding="utf-8") as f:
with open("./data.txt") as f:
    lines = f.readlines()

relation_dict = {}

def extract_relations_num(line):
    relations = []
    temp = line.split("mso:")
    if line.count("mso:") == 1:
        return 1
    elif line.count("mso:") == 2:
        return 2
    elif line.count("mso:") == 3:
        return 3
    return 0

train_data = []

for i,line in enumerate(lines):
    if i%3==1:
        relations_num = extract_relations_num(line)
        one_pair = []
        one_pair.append(extract_id(line))
        one_pair.append(extract_question(lines[i - 1]))
        one_pair.append(relations_num)
        train_data.append(one_pair)


print("start to load word dict!")


word_dict = {}
with open("word_dict.txt","r") as f:
    lines2 = f.readlines()
    for line in lines2:
        temp = line.split(" ")
        word_dict[temp[0]] = int(temp[1].replace("\n",""))



for item in train_data:
    new_sentence = replace_entity_in_logic_form(item[1],lines[int(item[0])*3-2])
    item[1] = new_sentence

with open("./train_nlpcc4.csv",mode="w") as f,open("./test_nlpcc4.csv",mode="w") as f2:
    f.write("id")
    f.write(",")
    f.write("text")
    f.write(",")
    f.write("one")
    f.write(",")
    f.write("two")
    f.write(",")
    f.write("three")
    f.write(",")
    f.write("\n")

    f2.write("id")
    f2.write(",")
    f2.write("text")
    f2.write(",")
    f2.write("one")
    f2.write(",")
    f2.write("two")
    f2.write(",")
    f2.write("three")
    f2.write(",")
    f2.write("\n")
    for index,item in enumerate(train_data):
        if random.randint(1,10)!=1:
            f.write("-1")
            f.write(",")
            f.write(item[1])
            f.write(",")
            temp = [0]*3
            temp[item[2] - 1]=1
            for i in temp:
                f.write(str(i))
                f.write(",")
            f.write("\n")
        else:
            f2.write("-1")
            f2.write(",")
            f2.write(item[1])
            f2.write(",")
            temp = [0] * 3
            temp[item[2] - 1] = 1
            for i in temp:
                f2.write(str(i))
                f2.write(",")
            f2.write("\n")


